# -*- coding: utf-8 -*-
"""
数据查询脚本 - 从JSON文件读取数据并查询知识库，使用大模型分析结果
"""

import json
import requests
from knowledge import KnowledgeBase
from config import OLLAMA_CONFIG

def read_json_file(file_path):
    """读取JSON文件"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"✅ 成功读取JSON文件: {file_path}")
        return data
    except Exception as e:
        print(f"❌ 读取JSON文件失败: {e}")
        return None

def analyze_with_llm(key, value, search_results):
    """
    使用大模型分析检索结果，提取最符合key的value值
    
    Args:
        key: 查询的键名（如"城市"）
        value: 原始值（如"辽源市"）
        search_results: 知识库检索结果
    
    Returns:
        dict: 分析结果，包含提取的值和推理过程
    """
    if not search_results:
        return {
            "extracted_value": None,
            "confidence": 0.0,
            "reasoning": "未找到相关信息"
        }
    
    # 构建检索内容文本
    context_texts = []
    for i, result in enumerate(search_results[:3]):  # 最多使用前3个结果
        context_texts.append(f"文档{i+1}: {result['text']}")
    
    context = "\n".join(context_texts)
    
    # 构建提示词
    prompt = f"""
基于以下检索到的相关文档，分析并提取关于"{key}"的具体值。

原始查询：{key} = {value}

相关文档：
{context}

请仔细分析这些文档，回答以下问题：
1. 从文档中能提取出关于"{key}"的具体值是什么？
2. 这个值与原始值"{value}"是否一致？
3. 你对这个提取结果的置信度如何（0-1）？

请按以下JSON格式回答：
{{
    "extracted_value": "提取的具体值",
    "confidence": 0.95,
    "reasoning": "基于文档X，可以确定..."
}}
"""
    
    try:
        # 调用Ollama API
        response = requests.post(
            f"{OLLAMA_CONFIG['base_url']}/api/generate",
            json={
                "model": OLLAMA_CONFIG.get('chat_model', 'llama2'),
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": 0.1,  # 降低随机性
                    "top_p": 0.9
                }
            },
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            answer = result.get('response', '').strip()
            
            # 尝试解析JSON响应
            try:
                # 提取JSON部分
                start_idx = answer.find('{')
                end_idx = answer.rfind('}') + 1
                if start_idx != -1 and end_idx > start_idx:
                    json_str = answer[start_idx:end_idx]
                    analysis = json.loads(json_str)
                    return analysis
                else:
                    # 如果没有找到JSON，使用默认解析
                    return {
                        "extracted_value": value,  # 保持原值
                        "confidence": 0.5,
                        "reasoning": f"大模型回答：{answer[:200]}..."
                    }
            except json.JSONDecodeError:
                return {
                    "extracted_value": value,
                    "confidence": 0.3,
                    "reasoning": f"解析失败，大模型原始回答：{answer[:200]}..."
                }
        else:
            return {
                "extracted_value": None,
                "confidence": 0.0,
                "reasoning": f"API调用失败：{response.status_code}"
            }
            
    except Exception as e:
        return {
            "extracted_value": None,
            "confidence": 0.0,
            "reasoning": f"分析失败：{str(e)}"
        }

def query_knowledge_for_key(kb, key, value):
    """
    根据键值对查询知识库并进行大模型分析
    
    Args:
        kb: 知识库实例
        key: 键名
        value: 键值
    
    Returns:
        dict: 包含搜索结果和分析结果的字典
    """
    # 构建查询字符串，包含键名和键值
    query_text = f"{key} {value}"
    
    print(f"\n🔍 查询: {key} = {value}")
    
    try:
        # 在知识库中搜索相关信息
        results = kb.search_knowledge(query_text, limit=3, similarity_threshold=0.3)
        
        if results:
            print(f"  ✅ 找到 {len(results)} 条相关信息:")
            for i, result in enumerate(results, 1):
                print(f"    {i}. [相似度: {result['similarity']:.3f}]")
                print(f"       {result['text'][:100]}...")
                print()
            
            # 使用大模型分析结果
            print(f"  🤖 正在使用大模型分析...")
            analysis = analyze_with_llm(key, value, results)
            
            print(f"  📊 分析结果:")
            print(f"     提取值: {analysis['extracted_value']}")
            print(f"     置信度: {analysis['confidence']:.2f}")
            print(f"     推理: {analysis['reasoning'][:100]}...")
            
            return {
                "search_results": results,
                "analysis": analysis,
                "has_results": True
            }
        else:
            print("  ❌ 未找到相关信息")
            return {
                "search_results": [],
                "analysis": {
                    "extracted_value": None,
                    "confidence": 0.0,
                    "reasoning": "未找到相关信息"
                },
                "has_results": False
            }
    
    except Exception as e:
        print(f"  ❌ 查询失败: {e}")
        return {
            "search_results": [],
            "analysis": {
                "extracted_value": None,
                "confidence": 0.0,
                "reasoning": f"查询失败：{str(e)}"
            },
            "has_results": False
        }

def get_user_input(key, value):
    """获取用户输入的信息"""
    print(f"\n📝 请为 '{key}: {value}' 输入相关信息:")
    print("   (直接回车跳过，输入 'q' 退出)")
    
    user_input = input("   > ").strip()
    
    if user_input.lower() == 'q':
        return None
    elif user_input:
        return user_input
    else:
        return ""

def save_to_knowledge_base(kb, key, value, info):
    """将用户输入的信息保存到知识库"""
    try:
        text = f"{key}:{value} - {info}"
        metadata = {
            "type": "用户输入",
            "key": key,
            "value": value,
            "source": "query_data_script"
        }
        
        knowledge_id = kb.add_knowledge(text, metadata)
        print(f"  ✅ 信息已保存到知识库，ID: {knowledge_id}")
        return True
    
    except Exception as e:
        print(f"  ❌ 保存失败: {e}")
        return False

def save_results_to_json(original_data, query_results, user_inputs, output_file="analysis_results.json"):
    """将分析结果保存到JSON文件"""
    try:
        # 构建结果数据结构
        results = {
            "原始数据": original_data,
            "查询时间": "",
            "分析结果": {},
            "用户补充信息": user_inputs,
            "汇总": {}
        }
        
        # 添加时间戳
        from datetime import datetime
        results["查询时间"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # 处理每个键的查询结果
        extracted_values = {}
        high_confidence_count = 0
        
        for key, value in original_data.items():
            query_result = query_results.get(key, {})
            search_results = query_result.get("search_results", [])
            analysis = query_result.get("analysis", {})
            user_info = user_inputs.get(key, "")
            
            # 确定最终值
            final_value = value  # 默认保持原值
            if user_info:
                final_value = user_info
            elif analysis.get("extracted_value") and analysis.get("confidence", 0) > 0.7:
                final_value = analysis["extracted_value"]
            
            extracted_values[key] = final_value
            
            if analysis.get("confidence", 0) > 0.7:
                high_confidence_count += 1
            
            results["分析结果"][key] = {
                "原始值": value,
                "提取值": analysis.get("extracted_value"),
                "置信度": analysis.get("confidence", 0),
                "推理过程": analysis.get("reasoning", ""),
                "用户输入": user_info,
                "最终值": final_value,
                "知识库搜索结果": [
                    {
                        "内容": result['text'],
                        "相似度": result['similarity']
                    } for result in search_results
                ]
            }
        
        # 生成汇总信息
        results["汇总"] = {
            "总查询项目": len(original_data),
            "找到知识库信息": len([k for k, v in query_results.items() if v.get("has_results")]),
            "高置信度分析": high_confidence_count,
            "用户补充信息": len([k for k, v in user_inputs.items() if v]),
            "提取的值": extracted_values
        }
        
        # 保存到JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"💾 分析结果已保存到: {output_file}")
        return True
        
    except Exception as e:
        print(f"❌ 保存结果失败: {e}")
        return False

def main():
    """主函数"""
    print("🚀 智能数据查询分析脚本启动")
    print("=" * 50)
    
    # 1. 读取JSON文件
    json_file = "data.json"
    data = read_json_file(json_file)
    
    if not data:
        print("❌ 无法读取数据文件，程序退出")
        return
    
    print(f"\n📋 读取到的数据:")
    for key, value in data.items():
        print(f"  {key}: {value}")
    
    # 2. 连接知识库
    try:
        print("\n💾 知识库配置")
        table_name = 'ly'
        
        if not table_name:
            print("使用默认表名...")
            kb = KnowledgeBase()
        else:
            print(f"使用自定义表名: {table_name}")
            kb = KnowledgeBase(table_name=table_name)
            
        print("✅ 知识库连接成功")
    except Exception as e:
        print(f"\n❌ 知识库连接失败: {e}")
        return
    
    # 3. 查询并分析每个键值对
    print("\n🔍 开始智能查询分析...")
    print("=" * 30)
    
    query_results = {}
    missing_info = {}
    
    for key, value in data.items():
        # 查询知识库并进行大模型分析
        result = query_knowledge_for_key(kb, key, value)
        query_results[key] = result
        
        # 如果没有找到相关信息，标记为缺失
        if not result.get("has_results"):
            missing_info[key] = value
    
    # 4. 处理缺失的信息
    user_inputs = {}
    
    if missing_info:
        print("\n❓ 以下信息在知识库中未找到:")
        print("=" * 30)
        
        for key, value in missing_info.items():
            print(f"\n缺失: {key} = {value}")
            
            # 获取用户输入
            user_info = get_user_input(key, value)
            
            if user_info is None:  # 用户选择退出
                print("👋 用户退出")
                break
            elif user_info:  # 用户提供了信息
                user_inputs[key] = user_info
                print(f"  ✅ 已记录用户输入")
            else:  # 用户跳过
                print(f"  ⏭️ 跳过 {key}")
    
    # 5. 显示最终结果
    print("\n📊 智能分析结果汇总:")
    print("=" * 30)
    
    for key, value in data.items():
        result = query_results.get(key, {})
        analysis = result.get("analysis", {})
        user_info = user_inputs.get(key, "")
        
        print(f"\n{key}: {value}")
        
        if result.get("has_results"):
            extracted_value = analysis.get("extracted_value")
            confidence = analysis.get("confidence", 0)
            
            print(f"  🤖 AI提取值: {extracted_value}")
            print(f"  🎯 置信度: {confidence:.2f}")
            
            if confidence > 0.8:
                print(f"  ✅ 高置信度匹配")
            elif confidence > 0.5:
                print(f"  ⚠️ 中等置信度")
            else:
                print(f"  ❌ 低置信度")
        
        if user_info:
            print(f"  📝 用户补充: {user_info}")
        
        if not result.get("has_results") and not user_info:
            print("  ❌ 无任何信息")
    
    # 6. 保存分析结果
    print("\n💾 保存分析结果...")
    save_results_to_json(data, query_results, user_inputs)
    
    # 7. 关闭连接
    kb.close()
    print(f"\n🎉 智能分析完成！共分析了 {len(data)} 个数据项")

if __name__ == "__main__":
    main()